1937D - Pinball - CodeForces Solution

#include<bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>  
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
#define int long long
#define pb push_back
#define tcs int T;cin>>T;while(T--)
#define rep(s,e,i) for (int i = s; i <= e; ++i)
#define revp(s,e,i) for (int i = s; i >= e; --i)
#define all(x) x.begin(), x.end()
#define out(x) cout<<((x)?"YES\n":"NO\n")
template <typename T> using ordered_set = 
tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
template <typename T> using ordered_multiset = 
tree<T, null_type, less_equal<T>, rb_tree_tag, tree_order_statistics_node_update>;
template <typename T> using min_heap = priority_queue<T,vector<T>,greater<T>>;
long long binexp(int x,int y,int mod){x=x%mod;int a=x,b=1;
while(y!=0){if(y&1){b=(1ll*b*a)%mod;}a=(1ll*a*a)%mod;y=(y>>1);}return b;}

signed main(){tcs{
int n;cin>>n;int psum1[n+2],psum2[n+2],p1[n+2],p2[n+2];
string s;cin>>s;rep(0,n+1,i){p1[i]=0;p2[i]=0;psum1[i]=0;psum2[i]=0;}
rep(1,n,i){psum1[i]=psum1[i-1];if(s[i-1]=='>'){++psum1[i];p1[psum1[i]]=p1[psum1[i]-1]+i;}}
revp(n,1,i){psum2[i]=psum2[i+1];if(s[i-1]=='<'){++psum2[i];p2[psum2[i]]=p2[psum2[i]-1]+i;}}
rep(1,n,i){int ans=0;
if(s[i-1]=='>'){
if(psum2[i+1]<=psum1[i-1]){ans=(p2[psum2[i+1]]-p1[psum1[i-1]]+p1[psum1[i-1]-psum2[i+1]])*2+n-i+1;}
else{ans=(p2[psum2[i+1]]-p1[psum1[i-1]]-p2[psum2[i+1]-psum1[i-1]-1])*2-i;}
}
else{
if(psum1[i-1]<=psum2[i+1]){ans=(p2[psum2[i+1]]-p2[psum2[i+1]-psum1[i-1]]-p1[psum1[i-1]])*2+i;}
else{ans=(p2[psum2[i+1]]-p1[psum1[i-1]]+p1[psum1[i-1]-psum2[i+1]-1])*2+n+i+1;}
}
cout<<ans<<" ";
}
cout<<endl;
}}